

# !/usr/bin/env python
# coding:utf-8
"""
@version:  1.0.0
@author:  Happylin
@file: re_fetch.py
@time: 2024/5/15 9:25
"""
import os
import re
import pandas as pd
import pdfplumber
from dateutil.parser import parse

class InvoiceExtract():



    def extract_invoice_date(self, text):
        # Identify potential date-related keywords
        keywords = ["开票日期"]
        for keyword in keywords:
            # print(keyword, text)
            pattern = rf"{keyword}\s*[：:]\s*(.*?)\n"
            matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
            if matches:
                for match in matches:
                    print("match:", match)
                    try:
                        # date = parse(match)  # Attempt to parse the extracted date
                        standard_date = match.replace("年", "-").replace("月", "-").replace("日", "")
                        return standard_date
                    except ValueError:
                        continue  # Ignore invalid date formats
        return None
    def extract_invoice_header(self, pdf_path):
        invoice_header = {
            "电子发票名称": os.path.basename(pdf_path),
            "发票代码": None,
            "发票号码": None,
            "开票日期": None,
            "购买方": None,
            "购买方-纳税人识别号": None,
            "销售方": None,
            "销售方-纳税人识别号": None,
            "备注": None,
            "总金额": None,
            "总税额": None,
            "价税合计": None,
            "发票备注": None
        }
        invoice_line =   {
            "项目名称": None,
            "规格型号": None,
            "单位": None,
            "数量": None,
            "单价": None,
            "金额": None,
            "税率": None,
            "税额": None
        }
        with pdfplumber.open(pdf_path) as pdf:
            for page in pdf.pages:
                # 提取文本内容
                text = page.extract_text()
                # print("text:", text)
                """
                如果text为空，说明无法提取文本内容，将'发票备注'字段写入'无法提取文本内容'，然后继续下一个循环
                发票代码、发票号码、开票日期、总金额、总税额、服务名称等字段的提取方法是：
                在text中找到包含'发票代码'、'发票号码'、'开票日期'、'合计'、'小写'等关键字的字符串，然后在这个字符串之后找到数字字符，然后提取这个数字字符
                这些对增值税发票和普通发票都适用
                但是发票行的提取会有所不同，因为增值税发票和普通发票的格式不同，所以提取方法也不同
                """
                if not text:
                    invoice_header["发票备注"] = "无法提取文本内容"
                    continue
                # print("text:", type(text), text)
                # 提取发票代码
                if not invoice_header["发票代码"]:
                    code_match = re.search(r'发票代码[:：]?\s*(\d+)', text)
                    if code_match:
                        invoice_header["发票代码"] = code_match.group(1).strip()
                # 提取发票号码
                if not invoice_header["发票号码"]:
                    number_match = re.search(r'发票号码[:：]?\s*(\d+)', text)
                    if number_match:
                        invoice_header["发票号码"] = number_match.group(1).strip()

                # 提取开票日期
                if not invoice_header["开票日期"]:
                    invoice_header["开票日期"] = self.extract_invoice_date(text)
                # 提取总金额    
                if not invoice_header["总金额"]:
                    total_amount_match = re.search(r'合\s*计\s*¥\s*(\d+\.\d+)', text)
                    if total_amount_match:
                        invoice_header["总金额"] = total_amount_match.group(1).strip()

                # 提取总税额
                # 提取方法是：在text中找到包含'合 计'的字符串，然后在这个字符串之后找到第二个'¥'字符，然后在这个字符之后找到数字字符，然后提取这个数字字符
                if not invoice_header["总税额"]:
                    total_tax_match = re.search(r'合\s*计.*?¥\s*\d+\.\d+.*?¥\s*(\d+\.\d+)', text, re.DOTALL)
                    if total_tax_match:
                        invoice_header["总税额"] = total_tax_match.group(1).strip()
                # 提取价税合计
                # 提取方法是：在text中找到包含'小写'的字符串，然后在这个字符串之后找到第一个'¥'字符，然后在这个字符之后找到数字字符，然后提取这个数字字符
                if not invoice_header["价税合计"]:
                    total_with_tax_match = re.search(r'小\s*写.*?¥\s*(\d+\.\d+)', text)
                    if total_with_tax_match:
                        invoice_header["价税合计"] = total_with_tax_match.group(1).strip()

                # 判断发票类型

                # 提取表格数据
                tables = page.extract_tables()
                columns = []
                if tables:
                    for row in tables[0]:
                        for cell in row:
                            # print("cell:", cell)
                            cells = []
                            if cell != None:
                                for item in cell.split('\n'):
                                    cells.append(item)
                            
                            columns.append(cells)       

                    # print("columns:",columns)
                    """
                    将columns的元素是列表，其中一个列表包含['购', '买', '方']， 
                    此列表之后的列表元素中是的第一个原始是一个包含'名称'和冒号的字符串，
                    将此字符串用冒号作为分隔符，冒号后面的字符串 作为'购买方'写到invoice_header中的'购买方'字段
                    """
                    """
                    在现在的程序逻辑基础上增加一个逻辑，逻辑为：与‘名称’同一个列表中的下一个元素的值是包含'纳税人识别号‘的值，
                    这个值用冒号分隔符分割后，第二个值作为’购买方-纳税人识别号写入到invoice_header中
                    """
                    # 在下面程序逻辑上再增加如下逻辑：在columns中找到包含'销'、'售'、'方'的列表元素，
                    # 然后在这个列表元素之后的元素中找到包含'名'和'称'两个字符的元素，这个元素用冒号分隔符分割后，
                    # 第二个值作为'销售方'写入到invoice_header中的'销售方'字段，此元素的下一个元素中包含'纳税人识别号'的值，
                    # 这个值用冒号分隔符分割后，第二个值作为'销售方-纳税人识别号'写入到invoice_header中的'销售方-纳税人识别号'字段
                    # 在上面程序逻辑上再增加如下逻辑：在columns中找到包含'货物或'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'项目名称'字段。然后继续在columns中找到包含'规格型号'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'规格型号'字段。
                    # 然后继续在columns中找到包含'单位'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'单位'字段。然后继续在columns中找到包含'数 量'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'数量'字段。
                    # 然后继续在columns中找到包含'单 价'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'单价'字段。然后继续在columns中找到包含'金 额'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'金额'字段。
                    # 然后继续在columns中找到包含'税率'字符的列表数据，刨除此列表中第一个元素，
                    # 后面的元素写入到invoice_line的'税率'字段。然后继续在columns中找到包含'税 额'字符的列表数据，
                    # 刨除此列表中第一个元素，后面的元素写入到invoice_line的'税额'字段。
                    for i, col in enumerate(columns):
                        print("col:", col, '购' in col)
                        if col:
                            if '货物' in col[0] or '服务' in col[0]:
                                invoice_line["项目名称"] = col[1:]
                            elif '规格型号' in col[0]:
                                invoice_line["规格型号"] = col[1:]
                            elif '单位' in col[0]:
                                invoice_line["单位"] = col[1:]
                            elif '数 量' in col[0]:
                                invoice_line["数量"] = col[1:]
                            elif '单 价' in col[0]:
                                invoice_line["单价"] = col[1:]
                            elif '金 额' in col[0]:
                                invoice_line["金额"] = col[1:]
                            elif '税率' in col[0]:
                                invoice_line["税率"] = col[1:]
                            elif '税 额' in col[0]:
                                invoice_line["税额"] = col[1:]   
                            elif '购' in col and '买' in col and '方' in col:
                                for subsequent_col in columns[i+1:]:
                                    if any('名' in item and '称' in item for item in subsequent_col):
                                        buyer_info = subsequent_col[0].split('：', 1)
                                        if len(buyer_info) > 1:
                                            invoice_header["购买方"] = buyer_info[1].strip()
                                        # Check for the next element containing '纳税人识别号'
                                        if len(subsequent_col) > 1:
                                            tax_id_info = subsequent_col[1].split('：', 1)
                                            if len(tax_id_info) > 1:
                                                invoice_header["购买方-纳税人识别号"] = tax_id_info[1].strip()
                                        break                                
                            elif '销' in col and '售' in col and '方' in col:
                                for subsequent_col in columns[i+1:]:
                                    if any('名' in item and '称' in item for item in subsequent_col):
                                        seller_info = subsequent_col[0].split('：', 1)
                                        if len(seller_info) > 1:
                                            invoice_header["销售方"] = seller_info[1].strip()
                                        # Check for the next element containing '纳税人识别号'
                                        if len(subsequent_col) > 1:
                                            tax_id_info = subsequent_col[1].split('：', 1)
                                            if len(tax_id_info) > 1:
                                                invoice_header["销售方-纳税人识别号"] = tax_id_info[1].strip()
                                        break                                

        return invoice_header, invoice_line


def extract_from_directory(directory_path):
    data = []
    invoices = {"Header": [], "Line": []}
    invoice_headers = []
    invoice_line = []
    for filename in os.listdir(directory_path):
        if filename.endswith(".pdf"):
            pdf_path = os.path.join(directory_path, filename)
            invoice_headers,invoice_line = InvoiceExtract().extract_invoice_header(pdf_path)
            invoices["Header"].append(invoice_headers)
            invoices["Line"].append(invoice_line)
    print("invoices:", invoices)        
    return invoices


if __name__ == '__main__':
    #print('PyCharm')
    extract_from_directory(r'D:\code\python\pythonProject4')

